"""
信息检索工作流

使用LangGraph的StateGraph实现信息检索流程：
- 信息检索
- 格式化结果
- 使用自定义事件流与债券分析工作流保持一致
- 已移除冗余的意图识别（由主工作流负责）
"""

import logging
from typing import TypedDict, Annotated, AsyncIterator, Dict, Any, List, Optional
from datetime import datetime
import operator

from langgraph.graph import StateGraph, END
from langgraph.checkpoint.memory import MemorySaver
from langgraph.config import get_stream_writer
from langchain_core.runnables import RunnableConfig

from agents.information_retrieval_agent import InformationRetrievalAgent


class RetrievalState(TypedDict):
    """信息检索状态"""
    # 输入
    user_query: str
    search_params: Dict[str, Any]  # 检索参数（关键词、评级、价格区间等）
    
    # ✅ 检索结果改为单个字符串（避免累积）
    retrieval_result: str
    
    # 元数据
    query_type: str  # basic_info, market_data, news, rating, performance
    errors: Annotated[List[str], operator.add]
    completed_steps: Annotated[List[str], operator.add]


class InformationRetrievalWorkflow:
    """信息检索工作流"""
    
    def __init__(self):
        """初始化工作流"""
        self.logger = logging.getLogger("workflows.information_retrieval")
        
        # 初始化智能体
        self.retrieval_agent = InformationRetrievalAgent()
        
        # ✅ 移除检查点保存器，避免状态累积
        self.checkpointer = None
        
        # 构建工作流图
        self.app = self._build_workflow()
    
    def _build_workflow(self):
        """构建工作流状态图（已移除冗余的意图识别和格式化节点）"""
        # 创建状态图
        workflow = StateGraph(RetrievalState)
        
        # 添加节点
        workflow.add_node("retrieval", self._retrieval_node)
        
        # ✅ 直接从入口点执行检索到结束（格式化由UI层负责）
        workflow.set_entry_point("retrieval")
        workflow.add_edge("retrieval", END)
        
        # 编译工作流（不使用checkpointer，确保每次检索独立）
        return workflow.compile()
    
    async def _retrieval_node(self, state: RetrievalState) -> Dict[str, Any]:
        """信息检索节点 - 使用自定义事件流"""
        try:
            user_query = state.get("user_query", "")
            search_params = state.get("search_params", {})
            
            # ✅ 获取stream writer
            writer = get_stream_writer()
            
            self.logger.info(f"开始信息检索: {user_query}")
            
            # 发送节点开始事件
            writer({
                "type": "step_start",
                "node": "retrieval",
                "step_name": "🔍 信息检索",
                "timestamp": datetime.now().isoformat()
            })
            
            # ✅ 直接传递user_query字符串，智能体会自动解析债券代码
            results = []
            
            # ✅ 流式执行智能体，转发所有推理事件
            async for event in self.retrieval_agent.stream_analyze(user_query):
                event_type = event.get("type")
                
                if event_type == "reasoning":
                    # ✅ 转发所有推理内容
                    content = event.get("content", "")
                    if content:
                        writer({
                            "type": "reasoning",
                            "node": "retrieval",
                            "content": content,
                            "timestamp": datetime.now().isoformat()
                        })
                
                elif event_type == "final_answer":
                    # 只保存结果，不发送到UI（避免在Step中显示完整结果）
                    results.append(event.get("content", ""))
            
            # ✅ 拼接为单个字符串
            retrieval_result = "\n".join(results) if results else "检索结果为空"
            
            self.logger.info(f"检索完成，结果数量: {len(results)}")
            
            # 发送节点结束事件
            writer({
                "type": "step_end",
                "node": "retrieval",
                "step_name": "🔍 信息检索",
                "timestamp": datetime.now().isoformat()
            })
            
            # ✅ 构建完整状态
            complete_state = {
                "user_query": user_query,
                "search_params": search_params,
                "retrieval_result": retrieval_result,  # ✅ 字符串
                "query_type": state.get("query_type", "basic_info"),
                "errors": state.get("errors", []),
                "completed_steps": state.get("completed_steps", []) + ["retrieval"]
            }
            
            # ✅ 发送final_result事件
            writer({
                "type": "final_result",
                "state": complete_state,
                "timestamp": datetime.now().isoformat()
            })
            
            return {
                "retrieval_result": retrieval_result,  # ✅ 字符串
                "completed_steps": ["retrieval"]
            }
        except Exception as e:
            self.logger.error(f"信息检索失败: {e}", exc_info=True)
            writer = get_stream_writer()
            writer({
                "type": "error",
                "node": "retrieval",
                "error": str(e),
                "timestamp": datetime.now().isoformat()
            })
            return {
                "retrieval_result": f"信息检索失败: {str(e)}",  # ✅ 字符串
                "errors": [f"信息检索失败: {str(e)}"],
                "completed_steps": ["retrieval"]
            }
    
    
    async def execute(
        self, 
        user_query: str,
        search_params: Optional[Dict[str, Any]] = None,
        thread_id: Optional[str] = None
    ) -> Dict[str, Any]:
        """
        执行信息检索工作流
        
        Args:
            user_query: 用户查询
            search_params: 检索参数（可选）
            thread_id: 线程ID（用于状态持久化）
            
        Returns:
            检索结果
        """
        if not thread_id:
            thread_id = f"info_retrieval_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
        
        # 初始化状态
        initial_state: RetrievalState = {
            "user_query": user_query,
            "search_params": search_params or {},
            "retrieval_result": "",  # ✅ 空字符串
            "query_type": "basic_info",
            "errors": [],
            "completed_steps": []
        }
        
        try:
            # 执行工作流
            config = RunnableConfig(configurable={"thread_id": thread_id})
            final_state = await self.app.ainvoke(initial_state, config)
            
            # 格式化结果
            return {
                "status": "success",
                "user_query": user_query,
                "timestamp": datetime.now().isoformat(),
                "results": {
                    "retrieval_data": final_state.get("retrieval_result", "")  # ✅ 直接使用字符串
                },
                "metadata": {
                    "query_type": final_state.get("query_type", ""),
                    "completed_steps": final_state.get("completed_steps", []),
                    "errors": final_state.get("errors", [])
                }
            }
        except Exception as e:
            self.logger.error(f"工作流执行失败: {e}")
            return {
                "status": "error",
                "user_query": user_query,
                "error": str(e),
                "timestamp": datetime.now().isoformat()
            }
    
    async def stream_execute(
        self,
        user_query: str,
        search_params: Optional[Dict[str, Any]] = None,
        thread_id: Optional[str] = None
    ) -> AsyncIterator[Dict[str, Any]]:
        """
        流式执行信息检索工作流（使用LangGraph自定义事件流）
        
        ✅ 与债券分析工作流保持一致：
        - 使用get_stream_writer()在节点内发送自定义事件
        - 使用stream_mode="custom"接收实时推理内容
        
        Args:
            user_query: 用户查询
            search_params: 检索参数
            thread_id: 线程ID
            
        Yields:
            工作流执行的每个步骤事件，包括智能体的实时推理内容
        """
        if not thread_id:
            thread_id = f"info_retrieval_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
        
        # 初始化状态
        initial_state: RetrievalState = {
            "user_query": user_query,
            "search_params": search_params or {},
            "retrieval_result": "",  # ✅ 空字符串
            "query_type": "basic_info",
            "errors": [],
            "completed_steps": []
        }
        
        config = RunnableConfig(configurable={"thread_id": thread_id})
        
        try:
            # ✅ 使用stream_mode="custom"接收节点内通过writer发送的自定义事件
            async for event in self.app.astream(
                initial_state,
                config,
                stream_mode="custom"  # 关键：只接收自定义事件
            ):
                # 直接转发自定义事件给UI层
                yield event
            
        except Exception as e:
            self.logger.error(f"流式执行失败: {e}", exc_info=True)
            yield {
                "type": "error",
                "error": str(e),
                "timestamp": datetime.now().isoformat()
            }


# 全局实例
information_retrieval_workflow = InformationRetrievalWorkflow()


# 便捷函数
async def search_bonds(user_query: str, search_params: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
    """
    检索债券信息的便捷函数
    
    Args:
        user_query: 用户查询
        search_params: 检索参数
        
    Returns:
        检索结果
    """
    return await information_retrieval_workflow.execute(user_query, search_params)